# -*- coding: utf-8 -*-
"""
基于Web界面的银行客户查询助手
"""
import os
import asyncio
from typing import Optional
import dashscope
from qwen_agent.agents import Assistant
from qwen_agent.gui import WebUI
import pandas as pd
from sqlalchemy import create_engine
from qwen_agent.tools.base import BaseTool, register_tool
import matplotlib.pyplot as plt
import io
import base64
import time
import numpy as np

# 解决中文显示问题
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'SimSun', 'Arial Unicode MS']  # 优先使用的中文字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 定义资源文件根目录
ROOT_RESOURCE = os.path.join(os.path.dirname(__file__), 'resource')

# 配置数据库连接
MYSQL_USER = os.getenv('MYSQL_USER', 'root')
MYSQL_PASSWORD = os.getenv('MYSQL_PASSWORD', '')
MYSQL_HOST = 'localhost'
MYSQL_PORT = '3306'
DATABASE_NAME = 'bank'  # 数据库名更改为 bank

def create_connection():
    """创建数据库连接"""
    try:
        engine = create_engine(
            f'mysql+mysqlconnector://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_HOST}:{MYSQL_PORT}/{DATABASE_NAME}?charset=utf8mb4',
            connect_args={'connect_timeout': 10}, pool_size=10, max_overflow=20
        )
        return engine
    except Exception as e:
        print(f"数据库连接失败: {str(e)}")
        return None

# ====== 更新助手的 system_prompt，修改名称为百万客群经营助手 ======
system_prompt = """
您好！欢迎使用百万客群经营助手。

我是您的智能助手，专注于帮助您高效管理和分析高净值客户数据。以下是我的核心功能：
- 查询客户基本信息（姓名、年龄、职业等）
- 分析客户的资产分布（存款、理财、基金等）
- 查看客户的交易行为和APP使用习惯
- 生成可视化图表以辅助决策

以下是数据库中的表结构说明：

客户基础信息表结构如下：
CREATE TABLE customer_base (
    customer_id VARCHAR(32) PRIMARY KEY COMMENT '客户ID',
    name VARCHAR(100) COMMENT '客户姓名',
    age INT COMMENT '年龄',
    gender VARCHAR(10) COMMENT '性别',
    occupation VARCHAR(100) COMMENT '职业',
    occupation_type VARCHAR(50) COMMENT '职业类型标签（如：企业高管/互联网从业者/私营业主）',
    monthly_income DECIMAL(12,2) COMMENT '月收入',
    open_account_date VARCHAR(10) COMMENT '开户日期',
    lifecycle_stage VARCHAR(50) COMMENT '客户生命周期',
    marriage_status VARCHAR(20) COMMENT '婚姻状态',
    city_level VARCHAR(20) COMMENT '城市等级（一线/二线城市）',
    branch_name VARCHAR(100) COMMENT '开户网点'
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='客户基础信息表';

客户行为资产表结构如下：
CREATE TABLE customer_behavior_assets (
    id VARCHAR(32) PRIMARY KEY COMMENT '主键ID',
    customer_id VARCHAR(32) COMMENT '客户ID（关联customer_base表）',
    total_assets DECIMAL(16,2) COMMENT '总资产',
    deposit_balance DECIMAL(16,2) COMMENT '存款余额',
    financial_balance DECIMAL(16,2) COMMENT '理财余额',
    fund_balance DECIMAL(16,2) COMMENT '基金余额',
    insurance_balance DECIMAL(16,2) COMMENT '保险余额',
    asset_level VARCHAR(20) COMMENT '资产分层（50万以下、50-80万、80-100万、100万+）',
    deposit_flag TINYINT COMMENT '是否持有存款（1是0否）',
    financial_flag TINYINT COMMENT '是否持有理财（1是0否）',
    fund_flag TINYINT COMMENT '是否持有基金（1是0否）',
    insurance_flag TINYINT COMMENT '是否持有保险（1是0否）',
    product_count INT COMMENT '持有产品数量',
    financial_repurchase_count INT COMMENT '近1年理财复购次数',
    credit_card_monthly_expense DECIMAL(12,2) COMMENT '信用卡月均消费',
    investment_monthly_count INT COMMENT '月均投资交易次数',
    app_login_count INT COMMENT 'APP月均登录次数',
    app_financial_view_time INT COMMENT '理财页面月均停留时长(秒)',
    app_product_compare_count INT COMMENT '产品对比点击次数',
    last_app_login_time VARCHAR(19) NULL COMMENT '最近APP登录时间',
    last_contact_time VARCHAR(19) NULL COMMENT '最近联系时间',
    contact_result VARCHAR(50) COMMENT '联系结果',
    marketing_cool_period VARCHAR(10) COMMENT '营销冷却期（30天）',
    stat_month VARCHAR(7) COMMENT '统计月份（YYYY-MM）',
    FOREIGN KEY (customer_id) REFERENCES customer_base(customer_id),
    UNIQUE KEY uk_customer_month (customer_id, stat_month)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='客户行为资产表';

每当 exc_sql 工具返回 markdown 表格和图片时，你必须原样输出工具返回的全部内容（包括图片），不要只总结表格，也不要省略图片。这样用户才能直接看到表格和图片。
知识库中是以前做过的示例，如果有类似的需要进行逻辑上的参考。
"""

functions_desc = [
    {
        "name": "exc_sql",
        "description": "对于生成的SQL，进行SQL查询",
        "parameters": {
            "type": "object",
            "properties": {
                "sql_input": {
                    "type": "string",
                    "description": "生成的SQL语句",
                }
            },
            "required": ["sql_input"],
        },
    },
]

# ====== 会话隔离 DataFrame 存储 ======
# 用于存储每个会话的 DataFrame，避免多用户数据串扰
_last_df_dict = {}

def get_session_id(kwargs):
    """根据 kwargs 获取当前会话的唯一 session_id，这里用 messages 的 id"""
    messages = kwargs.get('messages')
    if messages is not None:
        return id(messages)
    return None
# ====== exc_sql 工具类实现 ======
@register_tool('exc_sql')
class ExcSQLTool(BaseTool):
    """
    SQL查询工具，执行传入的SQL语句并返回结果，并自动进行可视化。
    """
    description = '对于生成的SQL，进行SQL查询，并自动可视化'
    parameters = [{
        'name': 'sql_input',
        'type': 'string',
        'description': '生成的SQL语句',
        'required': True
    }, {
        'name': 'need_visualization',
        'type': 'boolean',
        'description': '是否需要可视化图表，默认为True',
        'required': False
    }]

    def call(self, params: str, **kwargs) -> str:
        import json
        args = json.loads(params)
        sql_input = args['sql_input']
        need_visualization = args.get('need_visualization', True)  # 默认需要可视化
        database = args.get('database', DATABASE_NAME)  # 使用配置的数据库名称
        engine = create_connection()
        
        try:
            df = pd.read_sql(sql_input, engine)
            
            # 检查是否是股票对比查询，如果是则不需要可视化
            if "对比" in sql_input or "比较" in sql_input:
                need_visualization = False
            
            # 生成Markdown表格，显示前5行和后5行
            if len(df) <= 10:
                md = df.to_markdown(index=False)
            else:
                # 显示前5行和后5行
                head_df = df.head(5)
                tail_df = df.tail(5)
                md = head_df.to_markdown(index=False) + "\n...\n" + tail_df.to_markdown(index=False)
            
            # 如果不需要可视化，直接返回表格数据
            if not need_visualization:
                # 如果是股票对比查询，添加涨跌幅计算公式和结果
                if "对比" in sql_input or "比较" in sql_input:
                    pct_chg_formula = "\n\n涨跌幅计算公式：\n涨跌幅(%) = (当日收盘价 - 昨日收盘价) / 昨日收盘价 × 100%\n\n"
                    if 'pct_chg' in df.columns:
                        pct_chg_formula += "查询结果中的涨跌幅(%)字段即为按上述公式计算的结果。\n"
                    return f"{md}{pct_chg_formula}"
                return md
            
            # 添加数据描述信息
            describe_md = "数据统计信息：\n" + df.describe().to_markdown()
            
            # 如果只有一行或没有数据，则不生成图表
            if len(df) <= 1:
                if len(df) > 0:
                    return f"{md}\n\n{describe_md}"
                return md
            
            # 自动创建目录
            save_dir = os.path.join(os.path.dirname(__file__), 'image_show')
            os.makedirs(save_dir, exist_ok=True)
            filename = f'bar_{int(time.time()*1000)}.png'
            save_path = os.path.join(save_dir, filename)
            img_path = os.path.join('image_show', filename)
            
            # 生成图表
            chart_path = generate_chart_png(df, save_path)
            # 根据数据量确定图表类型
            chart_type = "折线图" if len(df) > 20 else "柱状图"
            img_md = f'![{chart_type}]({img_path})'
            
            return f"{md}\n\n{describe_md}\n\n{img_md}"
        except Exception as e:
            return f"SQL执行或可视化出错: {str(e)}"
# ========== 通用可视化函数 ========== 
def generate_chart_png(df_sql, save_path):
    """生成图表并保存到文件"""
    try:
        columns = df_sql.columns
        x = np.arange(len(df_sql))
        
        # 获取object类型
        object_columns = df_sql.select_dtypes(include='O').columns.tolist()
        if columns[0] in object_columns:
            object_columns.remove(columns[0])
        
        num_columns = df_sql.select_dtypes(exclude='O').columns.tolist()
        
        plt.figure(figsize=(12, 6))
        
        # 判断是否使用折线图（当数据量大于20时）
        use_line_chart = len(df_sql) > 20
        
        # 如果有分类列，则绘制堆积柱状图
        if len(object_columns) > 0:
            # 对数据进行透视
            pivot_df = df_sql.pivot_table(index=columns[0], columns=object_columns, 
                                          values=num_columns, 
                                          fill_value=0)
            # 绘制堆积柱状图
            ax = plt.subplot()
            bottoms = None
            for col in pivot_df.columns:
                if use_line_chart:
                    ax.plot(pivot_df.index, pivot_df[col], label=str(col), marker='o', markersize=4)
                else:
                    ax.bar(pivot_df.index, pivot_df[col], bottom=bottoms, label=str(col))
                    if bottoms is None:
                        bottoms = pivot_df[col].copy()
                    else:
                        bottoms += pivot_df[col]
        else:
            # 根据数据量决定使用折线图还是柱状图
            if use_line_chart:
                # 绘制折线图
                for column in columns[1:]:
                    plt.plot(x, df_sql[column], label=column, marker='o', markersize=3)
                chart_type = "折线图"
            else:
                # 绘制普通柱状图
                bottom = np.zeros(len(df_sql))
                for column in columns[1:]:
                    plt.bar(x, df_sql[column], bottom=bottom, label=column)
                    bottom += df_sql[column]
                chart_type = "柱状图"
            
            # 处理 trade_date 字段的格式化
            if 'trade_date' in df_sql.columns:
                # 将 trade_date 转换为 datetime 类型
                df_sql['trade_date'] = pd.to_datetime(df_sql['trade_date'], format='%Y-%m-%d')
                # 格式化为 YYYY-MM-DD
                formatted_dates = df_sql['trade_date'].dt.strftime('%Y-%m-%d')
                
                # 设置横坐标标签，避免过密
                # 只显示部分关键日期标签
                step = max(1, len(df_sql) // 10)  # 每10个数据点显示一个标签
                x_ticks = np.arange(0, len(df_sql), step)
                x_labels = formatted_dates.iloc[::step].tolist()
                
                plt.xticks(x_ticks, x_labels, rotation=45)
            else:
                # 对于大量数据，也需要筛选横坐标标签
                if len(df_sql) > 20:
                    step = max(1, len(df_sql) // 10)
                    x_ticks = np.arange(0, len(df_sql), step)
                    x_labels = [str(df_sql[columns[0]].iloc[i]) for i in x_ticks]
                    plt.xticks(x_ticks, x_labels, rotation=45)
                else:
                    plt.xticks(x, df_sql[columns[0]], rotation=45)
            
        plt.legend()
        plt.title("资产数据统计")
        plt.xlabel(columns[0])
        plt.ylabel("数值")
        plt.tight_layout()
        plt.savefig(save_path)
        plt.close()
        
        return save_path
        
    except Exception as e:
        print(f"图表生成失败: {str(e)}")
        return None

# ====== 初始化银行客户助手服务 ======
def init_agent_service():
    """初始化银行客户助手服务"""
    llm_cfg = {
        'model': 'qwen-turbo',
        'timeout': 30,
        'retry_count': 3,
    }
    # ====== 更新 tools 配置，新增 tavily mcp 支持 ======
    tools = [
        {
            "mcpServers": {
                "tavily-mcp": {
                    # Tavily MCP配置
                    "command": "npx",
                    "args": ["-y", "tavily-mcp@0.1.4"],
                    "env": {
                        "TAVILY_API_KEY": os.getenv('TAVILY_API_KEY', ''),
                    },
                    "disabled": False,
                    "autoApprove": []
                },
            }
        },
        'exc_sql'
    ]
    try:
        bot = Assistant(
            llm=llm_cfg,
            name='百万客群经营助手',
            description='银行客户数据查询与分析',
            system_message=system_prompt,
            function_list= tools,#['exc_sql'],  # 确保工具已注册
           # files=['./faq.txt']
        )
        print("助手初始化成功！")
        return bot
    except Exception as e:
        print(f"助手初始化失败: {str(e)}")
        raise


# ====== 定义 app_gui 函数以启动 Web 界面 ======
def app_gui():
    """图形界面模式，提供 Web 图形界面"""
    try:
        print("正在启动 Web 界面...")
        # 初始化助手
        bot = init_agent_service()
        # 配置聊天界面，列举典型查询问题
        chatbot_config = {
            'prompt.suggestions': [
                '我行目前有多少客户？总资产管理规模是多少？',
                '客户的平均资产是多少？高净值客户的占比如何？',
                '客户年龄分布情况如何？哪个年龄段客户最多？',
                '客户持有的四类产品（存款、理财、基金、保险）分布情况如何？'
               
                
            ]
        }
        print("Web 界面准备就绪，正在启动服务...")
        # 启动 Web 界面
        WebUI(
            bot,
            chatbot_config=chatbot_config
        ).run()
    except Exception as e:
        print(f"启动 Web 界面失败: {str(e)}")
        print("请检查网络连接和 API Key 配置")


if __name__ == '__main__':
    app_gui()  # 图形界面模式（默认）